import os
import time
import warnings
import logging
from pathlib import Path

import numpy as np
import joblib
from sklearn.model_selection import train_test_split

from util import createXY

# LazyPredict
from lazypredict.Supervised import LazyClassifier

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
warnings.filterwarnings("ignore")


def ensure_dir(p: str):
    Path(p).mkdir(parents=True, exist_ok=True)


def main():
    logging.info("读取图像，生成 X 和 y（flat + CPU）")
    cache_dir = '.cache_flat'
    ensure_dir(cache_dir)
    X, y = createXY(train_folder=os.path.join('data', 'train'), dest_folder=cache_dir, method='flat')
    X = np.asarray(X, dtype=np.float32)
    y = np.asarray(y)

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2025, stratify=y)

    # 仅挑选一组常用且稳定的模型，控制训练时间
    include = [
        'KNeighborsClassifier',
        'LogisticRegression',
        'RandomForestClassifier',
        'LinearSVC',
        'SVC',
        'RidgeClassifier',
        'SGDClassifier',
        'ExtraTreesClassifier',
        'DecisionTreeClassifier',
        'GradientBoostingClassifier',
        'AdaBoostClassifier',
        'BaggingClassifier',
        'GaussianNB'
    ]

    clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None, include=include)

    t0 = time.time()
    models_df, predictions_df = clf.fit(X_train, X_test, y_train, y_test)
    elapsed = time.time() - t0
    logging.info("LazyClassifier 评估完成，用时 %.2f 秒", elapsed)

    # 保存排行榜
    ensure_dir('models')
    leaderboard_path = os.path.join('models', 'lazy_leaderboard.csv')
    models_df.to_csv(leaderboard_path, encoding='utf-8-sig')
    logging.info("排行榜已保存: %s", leaderboard_path)

    # 选择最佳模型（按 Accuracy）
    best_name = models_df.sort_values(by='Accuracy', ascending=False).index[0]
    logging.info("最佳模型: %s", best_name)

    # 获取已训练好的模型；不同版本的 lazypredict 属性名不同，做多备份兜底
    trained = None
    for attr in ['trained_models', 'models', 'models_']:
        if hasattr(clf, attr):
            maybe = getattr(clf, attr)
            if isinstance(maybe, dict) and best_name in maybe:
                trained = maybe[best_name]
                break

    # 如果拿不到已训练模型，则简单重训一个同名默认配置模型
    if trained is None:
        from sklearn.utils import all_estimators
        name_to_cls = {name: cls for name, cls in all_estimators(type_filter='classifier')}
        if best_name in name_to_cls:
            trained = name_to_cls[best_name]()
            trained.fit(X_train, y_train)
        else:
            # 兜底：使用 KNN
            from sklearn.neighbors import KNeighborsClassifier
            best_name = 'KNeighborsClassifier'
            trained = KNeighborsClassifier(n_neighbors=5)
            trained.fit(X_train, y_train)

    save_path = os.path.join('models', 'lazy_best_model.joblib')
    joblib.dump({'model': trained, 'label_map': {0: 'cat', 1: 'dog'}, 'feature': 'flat', 'best_name': best_name}, save_path)
    logging.info("最佳模型已保存: %s | 模型: %s", save_path, best_name)


if __name__ == '__main__':
    main()
